﻿Imports System.Net
Imports System.Net.Sockets
Imports Utils.IO

Public Class Client
    Private mClient As TcpClient
    Private mThread As Threading.Thread
    Private mEncryptionProviderEx As System.Security.Cryptography.Rijndael
    Private mBuddy As Entities.Buddy
    Private mIsOnline As Boolean

    Public Property Buddy() As Entities.Buddy
        Get
            Return mBuddy
        End Get
        Friend Set(ByVal value As Entities.Buddy)
            mBuddy = value
        End Set
    End Property

    Public ReadOnly Property IsOnline() As Boolean
        Get
            Return mIsOnline
        End Get
    End Property
    Public ReadOnly Property IsDisconnecting() As Boolean
        Get
            Return mIsDisconnecting
        End Get
    End Property

    Public Sub StartPairUp(ByVal host As String, ByVal port As Integer, ByVal buddyNr As Integer, ByVal pinCode As String)
        mNetworkId = host & ":" & port
        Dim clsPairUpInfo As New PairUpInfo
        clsPairUpInfo.Host = host
        clsPairUpInfo.Port = port
        clsPairUpInfo.BuddyNr = buddyNr
        clsPairUpInfo.PinCode = pinCode

        mThread = New Threading.Thread(AddressOf Me.WorkerPairUp)
        mThread.Priority = Threading.ThreadPriority.BelowNormal
        mThread.IsBackground = True
        mThread.Start(clsPairUpInfo)
    End Sub
    Public Sub StartConnect(ByVal isAutomatedRetry As Boolean)
        mNetworkId = mBuddy.Host & ":" & mBuddy.Port
        mThread = New Threading.Thread(AddressOf Me.WorkerConnect)
        mThread.Priority = Threading.ThreadPriority.BelowNormal
        mThread.IsBackground = True
        mThread.Start(isAutomatedRetry)
    End Sub

    Private Sub WorkerConnect(ByVal isAutomatedRetry As Object)
        Dim blnIsAutomatedRetry As Boolean = False
        If TypeOf isAutomatedRetry Is Boolean Then blnIsAutomatedRetry = CBool(isAutomatedRetry)
        Try
            mClient = New TcpClient

            Manager.OnClientStatusChanged(Me)

            mClient.Connect(mBuddy.Host, mBuddy.Port)

            If Not NegotiateEncryption(False) Then Return

            Dim clsArgs As New CommandBlocks.Authenticate.Request
            clsArgs.UniqueId = NetNode.Settings.UniqueId
            clsArgs.PassCode = mBuddy.PassCode_LoginAtBuddy
            clsArgs.Email = NetNode.Settings.Email
            clsArgs.Name = NetNode.Settings.Nickname

            Dim strArgs As String = Utils.Serialization.ToXml(clsArgs)
            PostMessage("AUTHENTICATE", strArgs)

            Dim clsReplyCommand As MessageBlock = ReadMessageBlock(30) ' CommandBlock = ReadCommandBlock(30)
            If clsReplyCommand.Type <> "AUTHENTICATEREPLY" Then
                If Not blnIsAutomatedRetry Then Manager.OnClientLogEvent(Me, "Did not get a AUTHENTICATEREPLY!")
                Return
            End If
            If Not CommandBlocks.Authenticate.HandleReply(Me, clsReplyCommand) Then
                Return
            End If
            If Not blnIsAutomatedRetry Then Manager.OnClientLogEvent(Me, String.Format("Authenticated with buddy {0} complete", mBuddy.Name))

            'now, go into the command loop
            CommandLoop()
        Catch ex As Threading.ThreadAbortException
        Catch ex As Exception
            If Not blnIsAutomatedRetry Then Manager.OnClientException(Me, ex)
        Finally
            Close()
            Manager.OnClientStatusChanged(Me)
        End Try
    End Sub

    Private Sub WorkerPairUp(ByVal parameter As Object)
        Try
            Dim clsPairUpInfo As PairUpInfo = CType(parameter, PairUpInfo)

            mClient = New TcpClient
            mClient.Connect(clsPairUpInfo.Host, clsPairUpInfo.Port)
            Manager.OnClientStatusChanged(Me)

            If Not NegotiateEncryption(False) Then
                'Manager.OnClientLogEvent(Me, "Encryption negotiation failed")
                Return
            Else
                'Manager.OnClientLogEvent(Me, "Encryption negotiation succeeded")
            End If

            'now, send the pair up command block
            Dim clsArgs As New CommandBlocks.PairUp.Request
            clsArgs.BuddyId = clsPairUpInfo.BuddyNr
            clsArgs.PinCode = clsPairUpInfo.PinCode
            clsArgs.Host = NetNode.Settings.Host
            clsArgs.Port = NetNode.Settings.Port
            clsArgs.Name = NetNode.Settings.Nickname
            clsArgs.Email = NetNode.Settings.Email
            clsArgs.PassCode = Utils.Misc.GeneratePassCode()
            clsArgs.UniqueId = NetNode.Settings.UniqueId

            PostMessage("PAIRUP", Utils.Serialization.ToXml(clsArgs))

            Dim clsReplyCommand As MessageBlock = ReadMessageBlock(30) ' CommandBlock = ReadCommandBlock(30)
            If clsReplyCommand.Type <> "PAIRUPREPLY" Then
                'Manager.OnClientLogEvent(Me, "Did not get a PAIRUPREPLY!")
                Return
            End If
            If Not CommandBlocks.PairUp.HandleReply(Me, clsPairUpInfo, clsArgs, clsReplyCommand) Then
                Return
            End If
            Manager.OnClientLogEvent(Me, String.Format("Pairup with buddy {0} complete", mBuddy.Name))
            'now, go into the command loop
            CommandLoop()
        Catch ex As Threading.ThreadAbortException
        Catch ex As Exception
            Manager.OnClientException(Me, ex)
        Finally
            Close()
            Manager.OnClientStatusChanged(Me)
        End Try
    End Sub

    Private mIsDisconnecting As Boolean = False

    Private Sub CommandLoop()
        mIsOnline = True
        Manager.OnClientStatusChanged(Me)
        Manager.OnClientLogEvent(Me, "Going into command loop")
        Try
            While True
                Dim clsCmd As MessageBlock = ReadMessageBlock(-1)
                If clsCmd Is Nothing Then Exit While
                Select Case clsCmd.Type
                    Case "BYE"
                        mIsDisconnecting = True
                        Manager.OnClientStatusChanged(Me)
                        PostMessage("BYEBYE")
                        Exit While
                    Case "BYEBYE"
                        mIsDisconnecting = True
                        Manager.OnClientStatusChanged(Me)
                        Exit While
                    Case "TARGET", "TARGETREPLY"
                        OnTargetMessage(clsCmd)
                    Case Else
                        Manager.OnClientLogEvent(Me, String.Format("Command {0} ignored", clsCmd.Message))
                End Select
            End While
        Finally
            mIsOnline = False
        End Try
    End Sub

    Private Sub OnTargetMessage(ByVal message As MessageBlock)
        If message.Type = "TARGETREPLY" Then
            If mMessageHooks.ContainsKey(message.Id) Then
                mMessageReplies.Add(message.Id, message)
                mMessageHooks(message.Id).Set()
                mMessageHooks.Remove(message.Id)
                Return
            End If
        End If
        Manager.OnTargetMessage(Me, message)
    End Sub

    Private mMessageHooks As New Generic.Dictionary(Of Integer, Threading.ManualResetEvent)
    Private mMessageReplies As New Generic.Dictionary(Of Integer, MessageBlock)
    Private mNextMessageId As Integer = 0
    Private mNextMessageIdLock As New Object

    Private Function GetNextMessageId() As Integer
        SyncLock mNextMessageIdLock
            If mNextMessageId = 2 ^ 30 Then mNextMessageId = 0
            mNextMessageId += 1
            Return mNextMessageId
        End SyncLock
    End Function

    Private Function GetReplyHook(ByVal message As MessageBlock) As Threading.ManualResetEvent
        Dim clsHook As New Threading.ManualResetEvent(False)
        SyncLock mMessageHooks
            mMessageHooks.Add(message.Id, clsHook)
        End SyncLock
        Return clsHook
    End Function

    Friend Sub PostMessage(ByVal type As String)
        PostMessage(type, Nothing, Nothing)
    End Sub
    Friend Sub PostMessage(ByVal type As String, ByVal message As String)
        PostMessage(type, message, Nothing)
    End Sub
    Friend Sub PostMessage(ByVal type As String, ByVal data() As Byte)
        PostMessage(type, Nothing, data)
    End Sub
    Friend Sub PostMessage(ByVal type As String, ByVal message As String, ByVal data() As Byte)
        Dim clsMessage As New MessageBlock
        clsMessage.Id = GetNextMessageId()
        clsMessage.Type = type
        clsMessage.Message = message
        clsMessage.BinaryData = data
        SendMessageBlock(clsMessage)
    End Sub

    ''' <summary>
    ''' Posts message and returns the MessageId
    ''' </summary>
    ''' <param name="target"></param>
    ''' <param name="message"></param>
    ''' <param name="data"></param>
    ''' <returns>MessageId</returns>
    ''' <remarks></remarks>
    Public Function SendMessageToTargetAsync(ByVal target As String, ByVal message As String, ByVal data() As Byte) As Integer
        Dim clsMessage As New MessageBlock
        clsMessage.Id = GetNextMessageId()
        clsMessage.Type = "TARGET"
        clsMessage.Target = target
        clsMessage.Message = message
        clsMessage.BinaryData = data
        SendMessageBlock(clsMessage)
        Return clsMessage.Id
    End Function

    ''' <summary>
    ''' Sends message to target and waits for the reply
    ''' </summary>
    ''' <param name="target"></param>
    ''' <param name="message"></param>
    ''' <param name="data"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Function SendMessageToTarget(ByVal target As String, ByVal message As String, ByVal data() As Byte) As String
        Dim clsMessage As New MessageBlock
        clsMessage.Id = GetNextMessageId()
        clsMessage.Type = "TARGET"
        clsMessage.Target = target
        clsMessage.Message = message
        clsMessage.BinaryData = data
        Dim clsHook As Threading.ManualResetEvent = GetReplyHook(clsMessage)

        SendMessageBlock(clsMessage)

        clsHook.WaitOne(30000, False)
        If mMessageReplies.ContainsKey(clsMessage.Id) Then
            Dim clsRet As MessageBlock = mMessageReplies.Item(clsMessage.Id)
            mMessageReplies.Remove(clsMessage.Id)
            Return clsRet.Message
        Else
            Return Nothing
        End If
    End Function

    Private Sub SendMessageBlock(ByVal message As MessageBlock)
        Using clsMemStream As New System.IO.MemoryStream()
            Using clsCryptoStream As New System.Security.Cryptography.CryptoStream(clsMemStream, mEncryptionProviderEx.CreateEncryptor, System.Security.Cryptography.CryptoStreamMode.Write)
                'write blockversion
                BlockStreamWriter.Write(clsCryptoStream, CShort(2))

                BlockStreamWriter.Write(clsCryptoStream, message.Id)
                BlockStreamWriter.Write(clsCryptoStream, message.Type)
                BlockStreamWriter.Write(clsCryptoStream, message.Message)
                BlockStreamWriter.Write(clsCryptoStream, message.BinaryData)
                BlockStreamWriter.Write(clsCryptoStream, message.Target)
            End Using

            Dim bDataEncrypted() As Byte = clsMemStream.ToArray

            BlockStreamWriter.Write(mClient.GetStream, bDataEncrypted)

            Dim iPayLoad As Integer = message.GetPayloadSize ' message.Data.Length + message.BinaryData.Length ' Data.Length + Command.Length
            Dim iOverhead As Integer = 4 + bDataEncrypted.Length - iPayLoad
            Dim sWasteFactor As Double = 100 * iOverhead / iPayLoad
            'Manager.OnClientLogEvent(Me, String.Format("Command block ""{0}"" sent with {1} byte payload and {2} byte overhead ({3:0.0}% overhead)", message.Type, iPayLoad, iOverhead, sWasteFactor))
        End Using
    End Sub
    Private Function ReadMessageBlock(ByVal timeOut As Integer) As MessageBlock
        If timeOut = -1 Then
            mClient.GetStream.ReadTimeout = timeOut
        Else
            mClient.GetStream.ReadTimeout = timeOut * 1000
        End If

        Dim bEncryptedData() As Byte = BlockStreamReader.ReadBytes(mClient.GetStream)
        'Manager.OnClientLogEvent(Me, String.Format("Got {0} byte encrypted package", bEncryptedData.Length))

        Dim clsRet As New MessageBlock

        Dim clsIn As New System.IO.MemoryStream(bEncryptedData)
        Using clsCS As New System.Security.Cryptography.CryptoStream(clsIn, mEncryptionProviderEx.CreateDecryptor, System.Security.Cryptography.CryptoStreamMode.Read)
            Dim iVersion As Integer = BlockStreamReader.ReadInteger(clsCS)
            clsRet.Id = BlockStreamReader.ReadInteger(clsCS)
            clsRet.Type = BlockStreamReader.ReadString(clsCS)
            clsRet.Message = BlockStreamReader.ReadString(clsCS)
            clsRet.BinaryData = BlockStreamReader.ReadBytes(clsCS)
            clsRet.Target = BlockStreamReader.ReadString(clsCS)
        End Using

        Return clsRet
    End Function

    Private mNetworkId As String
    Public ReadOnly Property NetworkId() As String
        Get
            Return mNetworkId
        End Get
    End Property

    Friend Sub SendLine(ByVal text As String)
        'Utils.IO.BlockStreamWriter.Write(mClient.GetStream, text)
        Dim bBytes() As Byte = System.Text.Encoding.UTF8.GetBytes(text & vbCrLf)
        mClient.GetStream.Write(bBytes, 0, bBytes.Length)
    End Sub
    Friend Sub CloseAndAbort()
        Try
            mThread.Abort()
        Catch
        End Try
        Close()
    End Sub
    Friend Sub Close()
        If Not mIsDisconnecting Then
            Try
                mClient.Client.Close()
            Catch
            End Try
        End If
        Manager.RemoveClient(Me)
    End Sub
    Friend Sub Disconnect()
        If mIsOnline Then PostMessage("BYE") ' SendTextMessage("BYE", "")
    End Sub

    Friend Sub StartServer(ByVal client As TcpClient)
        mClient = client
        mNetworkId = client.Client.RemoteEndPoint.ToString

        mThread = New Threading.Thread(AddressOf Me.WorkerServer)
        mThread.Priority = Threading.ThreadPriority.BelowNormal
        mThread.IsBackground = True
        mThread.Start()
    End Sub

    Private Sub WorkerServer()
        Try
            Manager.OnClientStatusChanged(Me)

            'first, negotiate encryption
            If Not NegotiateEncryption(True) Then
                'Manager.OnClientLogEvent(Me, "Encryption negotiation failed")
                Return
            Else
                'Manager.OnClientLogEvent(Me, "Encryption negotiation succeeded")
            End If

            'now, validate that the client is a valid buddy
            Dim clsCmdBlock As MessageBlock = ReadMessageBlock(30) ' CommandBlock = ReadCommandBlock(30)
            'Manager.OnClientLogEvent(Me, String.Format("{0} command block received", clsCmdBlock.Type))
            Select Case clsCmdBlock.Type
                Case "PAIRUP"
                    If Not CommandBlocks.PairUp.HandleRequest(Me, clsCmdBlock) Then Return
                Case "AUTHENTICATE"
                    If Not CommandBlocks.Authenticate.HandleRequest(Me, clsCmdBlock) Then Return
                Case Else
                    'Manager.OnClientLogEvent(Me, "Invalid first command block. Needs to be PAIRUP or AUTHENTICATE")
                    Return
            End Select

            'now, go into the command loop
            CommandLoop()

        Catch ex As Threading.ThreadAbortException
        Catch ex As Exception
            Manager.OnClientException(Me, ex)
        Finally
            Close()
            Manager.OnClientStatusChanged(Me)
        End Try
    End Sub

    Private Function NegotiateEncryption(ByVal actAsServer As Boolean) As Boolean
        'Manager.OnClientLogEvent(Me, "Encryption negotiation started")
        If actAsServer Then
            'since we're the "server" beeing called, we start the conversation by sending
            'the caller our public RSA key
            Dim clsRSA As New System.Security.Cryptography.RSACryptoServiceProvider(1024)
            SendLine("RSAKEY:" & clsRSA.ToXmlString(False))
            'Manager.OnClientLogEvent(Me, "Public RSA Key sent")

            'now, wait for the symmetric key encoded using our public rsa key
            Dim strRinjKey As String = ReadLine(30)
            If Not strRinjKey.ToUpper.StartsWith("RIJNKEY:") Then
                'Manager.OnClientLogEvent(Me, "Did not get RINJKEY response")
                Return False
            End If
            'remove header
            strRinjKey = strRinjKey.Substring(8)

            'Manager.OnClientLogEvent(Me, "Received encrypted Rijn-key, decrypting...")
            'decrypt the rijnkey
            Dim aRinjKeyEncrypted() As Byte = System.Convert.FromBase64String(strRinjKey)
            Dim aRinjKey() As Byte = clsRSA.Decrypt(aRinjKeyEncrypted, True)

            'Manager.OnClientLogEvent(Me, "Rijn-key decrypted")

            'now, wait for the symmetric key encoded using our public rsa key
            Dim strRinjIV As String = ReadLine(30)
            If Not strRinjIV.ToUpper.StartsWith("RIJNIV:") Then
                'Manager.OnClientLogEvent(Me, "Did not get RINJIV response")
                Return False
            End If
            'remove header
            strRinjIV = strRinjIV.Substring(7)

            'Manager.OnClientLogEvent(Me, "Received encrypted Rijn-IV, decrypting...")
            'decrypt the rijnkey
            Dim aRinjIVEncrypted() As Byte = System.Convert.FromBase64String(strRinjIV)
            Dim aRinjIV() As Byte = clsRSA.Decrypt(aRinjIVEncrypted, True)

            'Manager.OnClientLogEvent(Me, "Rijn-IV decrypted")

            'initialize the symmetric encryption provider to be used in this session
            mEncryptionProviderEx = New System.Security.Cryptography.RijndaelManaged
            mEncryptionProviderEx.Key = aRinjKey
            mEncryptionProviderEx.IV = aRinjIV

            'Manager.OnClientLogEvent(Me, "Rijn encryptor/decryptor setup")

            SendLine("READY")
        Else
            'we're the "client"
            'Manager.OnClientLogEvent(Me, "Waiting for RSA key from server node")

            'read lines first line that cointains the RSAKEY
            Dim strLine As String = ReadLine(30)
            If Not strLine.ToUpper.StartsWith("RSAKEY:") Then
                'Manager.OnClientLogEvent(Me, "Did not receive RSA key from server node")
                Return False
            End If
            'remove header
            strLine = strLine.Substring(7)
            'initialize our rsa provider with the public key sent
            'Manager.OnClientLogEvent(Me, "RSA Key received")
            Dim clsRSA As New System.Security.Cryptography.RSACryptoServiceProvider
            clsRSA.FromXmlString(strLine)
            'Manager.OnClientLogEvent(Me, "RSA encryptor setup successfully")

            'now, create the symmetric key to be used in this session
            mEncryptionProviderEx = System.Security.Cryptography.Rijndael.Create ' New System.Security.Cryptography.RijndaelManaged
            mEncryptionProviderEx.GenerateKey()
            mEncryptionProviderEx.GenerateIV()

            'Manager.OnClientLogEvent(Me, "Rijn encryptor/decryptor setup successfully")

            'now, encrypt our symmetric key using the remote nodes public rsa key
            Dim bRijnKey As Byte() = mEncryptionProviderEx.Key
            Dim bKeyEncrypted As Byte() = clsRSA.Encrypt(bRijnKey, True)
            Dim sKeyEncrypted As String = Convert.ToBase64String(bKeyEncrypted)
            SendLine("RIJNKEY:" & sKeyEncrypted)
            'Manager.OnClientLogEvent(Me, "RSA-encrypted Rijn-key sent to the server node")

            Dim bRijnIV As Byte() = mEncryptionProviderEx.IV
            Dim bIVEncrypted As Byte() = clsRSA.Encrypt(bRijnIV, True)
            Dim sIVEncrypted As String = Convert.ToBase64String(bIVEncrypted)
            SendLine("RIJNIV:" & sIVEncrypted)
            'Manager.OnClientLogEvent(Me, "RSA-encrypted Rijn-IV sent to the server node")

            strLine = ReadLine(30)
            'response should be "READY"
            If strLine = "READY" Then
                'Manager.OnClientLogEvent(Me, "Got READY response from server node")
                Return True
            Else
                'Manager.OnClientLogEvent(Me, "Did not get the READY response from server node")
                Return False
            End If
        End If
        Return True
    End Function

    Friend Function ReadLine(ByVal timeout As Integer) As String
        If timeout = -1 Then
            mClient.GetStream.ReadTimeout = -1
        Else
            mClient.GetStream.ReadTimeout = timeout * 1000
        End If
        Dim bBuffer(0) As Byte
        Dim bLast As Byte
        Dim bReturn As New Generic.List(Of Byte)
        While True
            Dim iReadBytes As Integer = mClient.GetStream.Read(bBuffer, 0, 1)
            If iReadBytes <> 1 Then Return Nothing
            If bBuffer(0) = 10 AndAlso bLast = 13 Then
                Exit While
            ElseIf bLast <> 0 Then
                bReturn.Add(bLast)
            End If
            bLast = bBuffer(0)
        End While
        Return System.Text.Encoding.UTF8.GetString(bReturn.ToArray)
    End Function




End Class
